from loadDataset import Load2022Dataset
from torch.utils.data import Dataset
from typing import Tuple
import numpy as np

class Dataset2022ForMulti(Dataset):
    def __init__(self, channelsNum: Tuple, matPath: str) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self.tarData, self.noTarData = dataset.getDataNdarray()
        self.signalLen = self.tarData.shape[2]
        self.dataTarSize, self.dataNonTargetSize = dataset.getDataSize()
        self.dataTarArraySize, self.dataNonTargetArraySize = dataset.getDataArraySize()
        self.channelsNum = channelsNum
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        self.tarDataIndex = 0
        self.nonTarDataIndex = 0
        self.personSignalLen = len(channelsNum)
        self.dataSize = self.dataTarSize + self.dataNonTargetSize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        while count < self.dataTarSize:
            subjectA = self.tarData[self.tarDataArrayIndex]
            subjectB = self.tarData[(self.tarDataArrayIndex + 1) % self.dataTarArraySize]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []
                dataBList = []
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0:self.signalLen, i])
                for j in self.channelsNum:
                    dataBList.append(subjectB[j, 0:self.signalLen, i])
                self.data[count, 0:self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen:2 * self.personSignalLen, :] = np.array(dataBList)
                self.label[count] = 1
                count += 1
            self.tarDataArrayIndex += 1
        while count - self.dataTarSize < self.dataNonTargetSize:
            subjectA = self.noTarData[self.nonTarDataArrayIndex]
            subjectB = self.noTarData[(self.nonTarDataArrayIndex + 1) % self.dataNonTargetArraySize]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []
                dataBList = []
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0:self.signalLen, i])
                for j in self.channelsNum:
                    dataBList.append(subjectB[j, 0:self.signalLen, i])
                self.data[count, 0:self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen:2 * self.personSignalLen, :] = np.array(dataBList)
                self.label[count] = 0
                count += 1
            self.nonTarDataArrayIndex += 1
        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]
